#!/bin/env python
# 
# Stress test for CS 360 threaded messaging server
#       Daniel Zappala
#       2011

import optparse
import socket
import sys
import threading
import time

class RequesterThread(threading.Thread):
    """ Open a socket to the server and send a series of requests """
    def __init__(self, host, port, repetitions, name):
        """ arguments to initialize the requestor:

        host -- host name of the server
        port -- port number of the server
        repetitions -- number of time to repeat sending the messages
        name -- name of this thread
        """
        threading.Thread.__init__(self)
        # daemon threads will die if the main program exits
        threading.Thread.daemon = True
        self.host = host
        self.port = port
        self.repetitions = repetitions
        self.name = name
        self.open_socket()

    def open_socket(self):
        """  open a socket to the server """
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((self.host,self.port))
        except:
            error = sys.exc_info()[0]
            reason = sys.exc_info()[1]
            print "Failed to Connect"
            print error, reason
            self.sock = False

    def send_put(self,iteration):
        """ send a put message to the server """
        try:
            self.sock.send("put bob sub%d %d\n%s" % (iteration,10,'abcdefghij'))
            reply = self.sock.recv(1024)
            if reply != "OK\n":
                print "[put FAILED]", reply
                return False
            return True
        except:
            error = sys.exc_info()[0]
            reason = sys.exc_info()[1]
            print "[FAILED]", error, reason
            return False

    def send_list(self):
        """ send a list message to the server """
        try:
            self.sock.send("list bob\n")
            # get the main part of the reply
            reply = ""
            while reply.find("\n") == -1:
                reply += self.sock.recv(1024)
            fields = reply.split()
            if fields[0] != "list":
                print "[FAILED]", reply
                return False
            # collect the lines in the reply
            number = int(fields[1])
            extra = reply[reply.find("\n")+1:]
            while extra.count("\n") < number:
                extra += self.sock.recv(1024)
            print "%s Listed %d messages" % (self.name,number)
            return True
        except:
            error = sys.exc_info()[0]
            reason = sys.exc_info()[1]
            print "[FAILED]", error, reason
            return False

    def send_get(self,iteration):
        """ send a get message to the server """
        try:
            # skip first iteration since some students count from 0 and
            # others from 1
            if (iteration <= 1):
                return True
            self.sock.send("get bob %d\n" % (iteration - 1))
            # get the reply
            reply = ""
            while reply.find("\n") == -1:
                reply += self.sock.recv(1024)
            fields = reply.split()
            if fields[0] != "message":
                print "[FAILED]", reply
                return False
            # collect the body of the message
            subject = fields[1]
            length = int(fields[2])
            extra = reply[reply.find("\n")+1:]
            while len(extra) < length:
                extra += self.sock.recv(1024)
            return True
        except:
            error = sys.exc_info()[0]
            reason = sys.exc_info()[1]
            print "[FAILED]", error, reason
            return False

    def send_messages(self,iteration):
        """ send a series of messages to the server """
        # send put
        if not self.send_put(iteration):
            return False
        # send list
        if not self.send_list():
            return False
        # send get
        if not self.send_get(iteration):
            return False
        return True

    def run(self):
        """ run this thread for a certain number of iterations """
	if not self.sock:
            return
        iteration = 0
        while iteration < self.repetitions:
            iteration += 1
            success = self.send_messages(iteration)
            sys.stdout.flush()
            if not success:
                return
        self.sock.close()

class WorkloadGenerator:
    """ Generate a set of threads to make requests to the server """
    def __init__(self, hostname, port, num_threads, repetitions):
        """ arguments to initialize the class:
        
        hostname -- host name of the server
        port -- port of the server
        num_threads -- number of threads to start
        repetitions -- number of repetitions to run
        """
        self.threads = []
        self.hostname = hostname
        self.port = port
        for i in range(num_threads):
            name = "thread%d" % i
            self.threads.append(RequesterThread(hostname, port, repetitions, name))
		
    def run(self):
        """ run the workload generator """
        for thread in self.threads:
            thread.start()
        for thread in self.threads:
            thread.join(30)
            if thread.isAlive():
                print "Waited too long ... aborting"
                return
			
if __name__ == "__main__":
    # parse arguments
    parser = optparse.OptionParser(usage="%prog hostname:port -t [threads] -r [repetitions]", version="%prog 1.0")
    parser.add_option("-t", "--threads", dest="threads", type="int", default=10,
                      help=	"number of busy threads to test")
    parser.add_option("-r", "--repetitions", dest="repetitions", type="int", default=10,
                      help=	"number of repetitions for each thread")
    (options, args) = parser.parse_args()

    try:
        parts = args[0].split(':')
        hostname = parts[0]
        port = int(parts[1])
    except:
        parser.error("You must specifiy a host:port argument.")
    # print welcome message
    print "Host:        %s" % (hostname)
    print "Port:        %d" % (port)
    print "Threads:     %d" % (int(options.threads))
    print "Repetitions: %d" % (int(options.repetitions))
    print "--------------------------------------------------------"
        
    # launch generators
    generator = WorkloadGenerator(hostname, port, int(options.threads),
                                  int(options.repetitions))
    generator.run()

